

import os, random


def get_data(dst_filepath, img_dir, num=200):
    with open(dst_filepath, 'w') as f:
        l = os.listdir(img_dir)
        l = sorted(l)
        for class_name in l:
            class_txt = class_name.replace('_', ' ')
            class_dir = os.path.join(img_dir, class_name)
            imglists = os.listdir(class_dir)
            print(class_txt, len(imglists))
            # random.shuffle(imglists)
            # for imgname in imglists[:num]:
            #     f.write(
            #         class_name + '/' + imgname + ';' +  class_txt + '\n'
            #     )



if __name__ == "__main__":
    get_data(
        './data/style_train_art.txt',
        '/home/wangxin/data/style/wikiart',
        10000000
    )

    # get_data(
    #     './data/style_test_art.txt',
    #     '/home/wangxin/data/style/wikiart',
    #     20
    # )











